/**
 * Copyright 2009 by Nathaniel Harward.
 * 
 * This file is part of jtyrant-net.
 * 
 * jtyrant-net is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * jtyrant-net is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 * 
 * You should have received a copy of the GNU Lesser General Public License
 * along with jtyrant-net.  If not, see <http://www.gnu.org/licenses/>
 */
package us.harward.jtn;

import java.io.IOException;
import java.math.BigDecimal;
import java.math.MathContext;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.channels.ByteChannel;
import java.nio.channels.SocketChannel;
import java.util.Arrays;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

/**
 * Represents a connection to a remote Tokyo Tyrant server; one level up from
 * the socket itself. This class is a very thin wrapper around a {@link Socket}
 * object and provides no connection pooling, asynchronous I/O, etc. Calling
 * methods in this class translate directly into read/write calls on the
 * underlying socket.
 * 
 * All methods use the assert keyword to verify that their arguments are non-
 * <code>null</code> and that arguments meet certain requirements (see javadocs
 * for special requirements) -- disabling assertions will bypass these checks.
 * 
 * Please note that while this class can be shared by different threads to
 * invoke calls, it is explicitly <em>not</em> thread safe in the traditional
 * sense. Calls are synchronous and atomic, however two simultaneous calls from
 * different threads will interfere with each other. If a
 * {@link TokyoTyrantConnection} object is in a properly shared pool there
 * should not be any issues (similar to JDBC connections).
 * 
 * Please see the <a
 * href="http://tokyocabinet.sourceforge.net/tyrantdoc/#protocol">network
 * protocol description</a> for documentation on the method calls and their
 * arguments.
 */
public final class TokyoTyrantConnection {

    private static BigDecimal TRILLIONNUM = new BigDecimal(1000000000000L);

    private final ByteChannel channel;

    /**
     * Creates a new Tokyo Tyrant native binary network protocol wrapper around
     * an existing {@link ByteChannel}.
     * 
     * @param channel
     * @throws IOException
     * 
     * @throws AssertionError
     *             if <code>channel</code> is <code>null</code> or not connected
     *             (see {@link SocketChannel#isConnected()})
     */
    TokyoTyrantConnection(final ByteChannel channel) throws IOException {
	assert channel != null;
	this.channel = channel;
    }

    /**
     * Closes the connection to the remote server.
     * 
     * @throws IOException
     */
    public void close() throws IOException {
	channel.close();
    }

    public void put(final byte[] key, final byte[] value)
	    throws FailedCallException, IOException {
	if (parameterizedPut(Magic.put, key, value, null, true) != 0)
	    throw new FailedCallException(ErrorCode.EMISC, "put");
    }

    public void putkeep(final byte[] key, final byte[] value)
	    throws IOException, FailedCallException {
	if (parameterizedPut(Magic.putkeep, key, value, null, true) != 0)
	    throw new FailedCallException(ErrorCode.EKEEP, "putkeep");
    }

    public void putcat(final byte[] key, final byte[] value)
	    throws FailedCallException, IOException {
	if (parameterizedPut(Magic.putcat, key, value, null, true) != 0)
	    throw new FailedCallException(ErrorCode.EMISC, "putcat");
    }

    public void putshl(final byte[] key, final byte[] value, final int width)
	    throws IOException, FailedCallException {
	if (parameterizedPut(Magic.putshl, key, value, width, true) != 0)
	    throw new FailedCallException(ErrorCode.EMISC, "putcat");
    }

    public void putnr(final byte[] key, final byte[] value) throws IOException {
	parameterizedPut(Magic.putnr, key, value, null, false);
    }

    private int parameterizedPut(final ByteBuffer magic, final byte[] key,
	    final byte[] value, final Integer width,
	    final boolean readReturnCode) throws IOException {
	assert key != null;
	assert value != null;
	write(magic);
	writeInt32(key.length);
	writeInt32(value.length);
	if (width != null)
	    writeInt32(width);
	write(ByteBuffer.wrap(key), ByteBuffer.wrap(value));
	if (readReturnCode) {
	    return readUInt8();
	} else {
	    return 1;
	}
    }

    public void out(final byte[] key) throws FailedCallException, IOException {
	assert key != null;
	write(Magic.out);
	writeInt32(key.length);
	write(ByteBuffer.wrap(key));
	if (readUInt8() != 0)
	    throw new FailedCallException(ErrorCode.ENOREC, "out");
    }

    public byte[] get(final byte[] key) throws FailedCallException, IOException {
	assert key != null;
	write(Magic.get);
	writeInt32(key.length);
	write(ByteBuffer.wrap(key));
	if (readUInt8() != 0) {
	    throw new FailedCallException(ErrorCode.ENOREC, "get");
	} else {
	    final int size = readInt32();
	    final byte[] rv = new byte[size];
	    read(ByteBuffer.wrap(rv));
	    return rv;
	}
    }

    public Map<byte[], byte[]> mget(final byte[]... keys)
	    throws FailedCallException, IOException {
	return mget(Arrays.asList(keys));
    }

    public Map<byte[], byte[]> mget(final List<byte[]> keys)
	    throws FailedCallException, IOException {
	assert keys != null;
	write(Magic.mget);
	writeInt32(keys.size());
	for (byte[] key : keys) {
	    writeInt32(key.length);
	    write(ByteBuffer.wrap(key));
	}
	if (readUInt8() != 0) {
	    throw new FailedCallException(ErrorCode.ENOREC, "mget");
	} else {
	    final int rnum = readInt32();
	    final Map<byte[], byte[]> rv = new HashMap<byte[], byte[]>();
	    for (int pos = 0; pos < rnum; ++pos) {
		final byte[] key = new byte[readInt32()];
		final byte[] value = new byte[readInt32()];
		read(ByteBuffer.wrap(key));
		read(ByteBuffer.wrap(value));
		rv.put(key, value);
	    }
	    return rv;
	}
    }

    public int vsiz(final byte[] key) throws FailedCallException, IOException {
	assert key != null;
	write(Magic.vsiz);
	writeInt32(key.length);
	write(ByteBuffer.wrap(key));
	if (readUInt8() != 0) {
	    throw new FailedCallException(ErrorCode.ENOREC, "vsiz");
	} else {
	    return readInt32();
	}
    }

    public void iterinit() throws FailedCallException, IOException {
	write(Magic.iterinit);
	if (readUInt8() != 0)
	    throw new FailedCallException(ErrorCode.EMISC, "iterinit");
    }

    public byte[] iternext() throws FailedCallException, IOException {
	write(Magic.iternext);
	if (readUInt8() != 0) {
	    throw new FailedCallException(ErrorCode.ENOREC, "iternext");
	} else {
	    final byte[] rv = new byte[readInt32()];
	    read(ByteBuffer.wrap(rv));
	    return rv;
	}
    }

    /**
     * @param prefix
     *            the prefix to forward match on
     * @param maxKeys
     *            the maximum number of keys to return; negative values mean no
     *            limit
     * @return a {@link List} of keys
     * @throws FailedCallException
     * @throws IOException
     */
    public List<byte[]> fwmkeys(final byte[] prefix, final int maxKeys)
	    throws FailedCallException, IOException {
	assert prefix != null;
	write(Magic.fwmkeys);
	writeInt32(prefix.length);
	writeInt32(maxKeys);
	write(ByteBuffer.wrap(prefix));
	if (readUInt8() != 0) {
	    throw new FailedCallException(ErrorCode.ENOREC, "fwmkeys");
	} else {
	    final int rnum = readInt32();
	    final List<byte[]> rv = new LinkedList<byte[]>();
	    for (int pos = 0; pos < rnum; ++pos) {
		final byte[] key = new byte[readInt32()];
		read(ByteBuffer.wrap(key));
		rv.add(key);
	    }
	    return rv;
	}
    }

    public int addint(final byte[] key, final int value)
	    throws FailedCallException, IOException {
	assert key != null;
	write(Magic.addint);
	writeInt32(key.length);
	writeInt32(value);
	write(ByteBuffer.wrap(key));
	if (readUInt8() != 0) {
	    throw new FailedCallException(ErrorCode.EKEEP, "addint");
	} else {
	    return readInt32();
	}
    }

    /**
     * Wrapper around {@link #adddouble(byte[], BigDecimal)}.
     */
    public double adddouble(final byte[] key, final double value)
	    throws FailedCallException, IOException {
	assert key != null;
	return adddouble(key, new BigDecimal(value)).doubleValue();
    }

    public BigDecimal adddouble(final byte[] key, final BigDecimal value)
	    throws FailedCallException, IOException {
	assert key != null;
	assert value != null;
	write(Magic.adddouble);
	writeInt32(key.length);
	final BigDecimal[] parts = value.divideAndRemainder(new BigDecimal(1),
		MathContext.UNLIMITED);
	writeInt64(parts[0].longValue());
	writeInt64(parts[1].multiply(TRILLIONNUM, MathContext.UNLIMITED)
		.longValue());
	write(ByteBuffer.wrap(key));
	if (readUInt8() != 0) {
	    throw new FailedCallException(ErrorCode.EKEEP, "adddouble");
	} else {
	    final BigDecimal integral = new BigDecimal(readInt64());
	    final BigDecimal fractional = new BigDecimal(readInt64()).divide(
		    TRILLIONNUM, MathContext.UNLIMITED);
	    return integral.add(fractional, MathContext.UNLIMITED);
	}
    }

    public byte[] ext(final byte[] functionName, final int options,
	    final byte[] key, final byte[] value) throws IOException,
	    FailedCallException {
	assert functionName != null;
	assert key != null;
	assert value != null;
	write(Magic.ext);
	writeInt32(functionName.length);
	writeInt32(options);
	writeInt32(key.length);
	writeInt32(value.length);
	write(ByteBuffer.wrap(functionName), ByteBuffer.wrap(key), ByteBuffer
		.wrap(value));
	if (readUInt8() != 0) {
	    throw new FailedCallException(ErrorCode.EMISC, "ext");
	} else {
	    final byte[] rv = new byte[readInt32()];
	    read(ByteBuffer.wrap(rv));
	    return rv;
	}
    }

    public void sync() throws FailedCallException, IOException {
	write(Magic.sync);
	if (readUInt8() != 0)
	    throw new FailedCallException(ErrorCode.EMISC, "sync");
    }

    public void vanish() throws FailedCallException, IOException {
	write(Magic.vanish);
	if (readUInt8() != 0)
	    throw new FailedCallException(ErrorCode.EMISC, "vanish");
    }

    public void copy(final byte[] path) throws FailedCallException, IOException {
	assert path != null;
	write(Magic.copy);
	writeInt32(path.length);
	write(ByteBuffer.wrap(path));
	if (readUInt8() != 0)
	    throw new FailedCallException(ErrorCode.EMISC, "copy");
    }

    public void restore(final byte[] path, final long timestamp)
	    throws FailedCallException, IOException {
	assert path != null;
	write(Magic.restore);
	writeInt32(path.length);
	writeInt64(timestamp);
	write(ByteBuffer.wrap(path));
	if (readUInt8() != 0)
	    throw new FailedCallException(ErrorCode.EMISC, "restore");
    }

    public void setmst(final byte[] hostname, final int port)
	    throws FailedCallException, IOException {
	assert hostname != null;
	write(Magic.setmst);
	writeInt32(hostname.length);
	writeInt32(port);
	write(ByteBuffer.wrap(hostname));
	if (readUInt8() != 0)
	    throw new FailedCallException(ErrorCode.EMISC, "setmst");
    }

    public long rnum() throws FailedCallException, IOException {
	write(Magic.rnum);
	if (readUInt8() != 0)
	    throw new FailedCallException(ErrorCode.EMISC, "rnum");
	return readInt64();
    }

    public long size() throws FailedCallException, IOException {
	write(Magic.size);
	if (readUInt8() != 0)
	    throw new FailedCallException(ErrorCode.EMISC, "size");
	return readInt64();
    }

    public byte[] stat() throws FailedCallException, IOException {
	write(Magic.stat);
	if (readUInt8() != 0) {
	    throw new FailedCallException(ErrorCode.EMISC, "stat");
	} else {
	    final byte[] rv = new byte[readInt32()];
	    read(ByteBuffer.wrap(rv));
	    return rv;
	}
    }

    public List<byte[]> misc(final byte[] functionName, final int options,
	    final byte[]... arguments) throws IOException, FailedCallException {
	return misc(functionName, options, Arrays.asList(arguments));
    }

    public List<byte[]> misc(final byte[] functionName, final int options,
	    final List<byte[]> arguments) throws IOException,
	    FailedCallException {
	assert functionName != null;
	write(Magic.misc);
	writeInt32(functionName.length);
	writeInt32(options);
	writeInt32(arguments.size());
	write(ByteBuffer.wrap(functionName));
	for (byte[] argument : arguments) {
	    writeInt32(argument.length);
	    write(ByteBuffer.wrap(argument));
	}
	if (readUInt8() != 0) {
	    throw new FailedCallException(ErrorCode.EMISC, "misc");
	} else {
	    final int rnum = readInt32();
	    final List<byte[]> rv = new LinkedList<byte[]>();
	    for (int pos = 0; pos < rnum; ++pos) {
		final byte[] result = new byte[readInt32()];
		read(ByteBuffer.wrap(result));
		rv.add(result);
	    }
	    return rv;
	}
    }

    private int readUInt8() throws IOException {
	final ByteBuffer buf = ByteBuffer.allocate(4);
	buf.position(buf.position() + 3);
	read(buf);
	buf.flip();
	return buf.getInt();
    }

    private void writeInt32(final int i) throws IOException {
	final ByteBuffer buf = ByteBuffer.allocate(4);
	buf.putInt(i);
	buf.flip();
	write(buf);
    }

    private void writeInt64(final long l) throws IOException {
	final ByteBuffer buf = ByteBuffer.allocate(8);
	buf.putLong(l);
	buf.flip();
	write(buf);
    }

    private int readInt32() throws IOException {
	final ByteBuffer buf = ByteBuffer.allocate(4);
	read(buf);
	buf.flip();
	return buf.getInt();
    }

    private long readInt64() throws IOException {
	final ByteBuffer buf = ByteBuffer.allocate(8);
	read(buf);
	buf.flip();
	return buf.getLong();
    }

    private void read(final ByteBuffer... byteBuffers) throws IOException {
	read(Arrays.asList(byteBuffers));
    }

    private void read(final List<ByteBuffer> byteBuffers) throws IOException {
	assert byteBuffers != null;
	for (ByteBuffer buf : byteBuffers)
	    while (buf.hasRemaining())
		channel.read(buf);
    }

    private void write(final ByteBuffer... byteBuffers) throws IOException {
	write(Arrays.asList(byteBuffers));
    }

    private void write(final List<ByteBuffer> byteBuffers) throws IOException {
	assert byteBuffers != null;
	for (ByteBuffer byteBuffer : byteBuffers) {
	    ByteBuffer readOnlyBuf = byteBuffer.asReadOnlyBuffer();
	    while (readOnlyBuf.hasRemaining())
		channel.write(readOnlyBuf);
	}
    }

}
